import { beforeEach, describe, expect, it, vi } from 'vitest'

import { ChainEventTypes } from '@latitude-data/constants'
import { ChainError, RunErrorCodes } from '@latitude-data/constants/errors'
import { Providers } from '@latitude-data/constants'
import { LogSources, StreamEventTypes } from '../../constants'
import { publisher } from '../../events/publisher'
import { ProviderLogsRepository } from '../../repositories'
import { createProject, createTelemetryContext } from '../../tests/factories'
import { testConsumeStream } from '../../tests/helpers'
import { Ok, Result } from './../../lib/Result'
import { runDocumentAtCommit } from './index'

const mocks = {
  publish: vi.fn(),
  uuid: vi.fn(() => 'fake-document-log-uuid'),
  runAi: vi.fn(async () => {
    const fullStream = new ReadableStream({
      start(controller) {
        controller.close()
      },
    })

    return Result.ok({
      type: 'text',
      text: Promise.resolve('Fake AI generated text'),
      providerLog: Promise.resolve({ uuid: 'fake-provider-log-uuid' }),
      usage: Promise.resolve({
        promptTokens: 0,
        completionTokens: 0,
        totalTokens: 0,
      }),
      toolCalls: Promise.resolve([]),
      response: Promise.resolve({ messages: [] }),
      fullStream,
    })
  }),
}

const dummyDoc1Content = `
---
provider: openai
model: gpt-4o
---

<step>
  This is a test document
</step>
<step>
  With two steps
</step>
`

async function buildData({ doc1Content }: { doc1Content: string }) {
  const { workspace, project, documents, commit, user, providers } =
    await createProject({
      providers: [{ type: Providers.OpenAI, name: 'openai' }],
      documents: {
        doc1: doc1Content,
      },
    })

  const context = await createTelemetryContext({ workspace })

  return {
    context,
    workspace,
    document: documents[0]!,
    commit,
    user,
    project,
    provider: providers[0]!,
  }
}

const publisherSpy = vi.spyOn(
  await import('../../events/publisher').then((f) => f.publisher),
  'publishLater',
)
const aiSpy = vi.spyOn(await import('../ai'), 'ai')
const createDocumentLogSpy = vi.spyOn(
  await import('../documentLogs/create'),
  'createDocumentLog',
)

describe('runDocumentAtCommit', () => {
  beforeEach(() => {
    vi.clearAllMocks()
    // @ts-expect-error - we are mocking the function
    aiSpy.mockImplementation(mocks.runAi)
    // @ts-expect-error - we are mocking the function
    publisherSpy.mockImplementation(mocks.publishLater)
  })

  describe('with an existing provider key', () => {
    it('returns document resolvedContent', async () => {
      const { context, workspace, document, commit } = await buildData({
        doc1Content: dummyDoc1Content,
      })
      const { lastResponse, duration, resolvedContent } =
        await runDocumentAtCommit({
          context,
          workspace,
          document,
          commit,
          parameters: {},
          source: LogSources.API,
        }).then((r) => r.unwrap())

      await lastResponse
      await duration

      expect(resolvedContent.trim()).toEqual(`---
provider: openai
model: gpt-4o
---


<step>
  This is a test document
</step>
<step>
  With two steps
</step>`)
    })

    it('pass params to AI', async () => {
      const { context, workspace, document, commit, provider } =
        await buildData({
          doc1Content: dummyDoc1Content,
        })

      const { lastResponse } = await runDocumentAtCommit({
        context,
        workspace,
        document,
        commit,
        parameters: {},
        source: LogSources.API,
      }).then((r) => r.unwrap())

      await lastResponse // wait for the chain to finish

      expect(aiSpy).toHaveBeenCalledWith(
        expect.objectContaining({
          messages: [
            {
              role: 'system',
              content: [
                {
                  type: 'text',
                  text: 'This is a test document',
                  _promptlSourceMap: [],
                },
              ],
            },
          ],
          config: { model: 'gpt-4o', provider: 'openai', maxSteps: 20 },
          provider,
        }),
      )
    })

    it('send documentLogUuid when chain is completed', async () => {
      const { context, workspace, document, commit } = await buildData({
        doc1Content: dummyDoc1Content,
      })
      const { lastResponse, duration, stream } = await runDocumentAtCommit({
        context,
        workspace,
        document,
        commit,
        parameters: {},
        source: LogSources.API,
      }).then((r) => r.unwrap())

      await lastResponse
      await duration

      const { value } = await testConsumeStream(stream)
      const repo = new ProviderLogsRepository(workspace.id)
      const logs = await repo.findAll().then((r) => r.unwrap())

      expect(value).toEqual(
        expect.arrayContaining([
          {
            event: StreamEventTypes.Latitude,
            data: expect.objectContaining({
              type: ChainEventTypes.ChainStarted,
            }),
          },
          {
            event: StreamEventTypes.Latitude,
            data: expect.objectContaining({
              type: ChainEventTypes.ProviderCompleted,
              providerLogUuid: logs[0]!.uuid,
            }),
          },
          {
            event: StreamEventTypes.Latitude,
            data: expect.objectContaining({
              type: ChainEventTypes.ProviderCompleted,
              providerLogUuid: logs[0]!.uuid,
            }),
          },
          {
            event: StreamEventTypes.Latitude,
            data: expect.objectContaining({
              type: ChainEventTypes.ProviderCompleted,
              providerLogUuid: logs[1]!.uuid,
            }),
          },
        ]),
      )

      expect(value.at(-1)).toEqual({
        event: StreamEventTypes.Latitude,
        data: expect.objectContaining({
          type: ChainEventTypes.ChainCompleted,
          uuid: logs[0]!.documentLogUuid,
          messages: expect.any(Array),
        }),
      })
    })

    it('calls publisher with correct data', async () => {
      const { context, workspace, document, commit } = await buildData({
        doc1Content: dummyDoc1Content,
      })
      const parameters = { testParam: 'testValue' }
      const { lastResponse, duration } = await runDocumentAtCommit({
        context,
        workspace,
        document,
        commit,
        parameters,
        source: LogSources.API,
      }).then((r) => r.unwrap())

      // Wait for the lastResponse promise to resolve
      await lastResponse
      await duration

      expect(publisher.publishLater).toHaveBeenCalled()
    })

    it('creates a document log', async () => {
      const { context, workspace, document, commit } = await buildData({
        doc1Content: dummyDoc1Content,
      })
      const { lastResponse } = await runDocumentAtCommit({
        context,
        workspace,
        document,
        commit,
        parameters: {},
        source: LogSources.API,
      }).then((r) => r.unwrap())

      await lastResponse

      expect(createDocumentLogSpy).toHaveResolvedWith(expect.any(Ok))
    })

    it('creates a document log with custom identifier', async () => {
      const { context, workspace, document, commit } = await buildData({
        doc1Content: dummyDoc1Content,
      })
      const parameters = { testParam: 'testValue' }
      const { lastResponse } = await runDocumentAtCommit({
        context,
        workspace,
        document,
        commit,
        parameters,
        customIdentifier: 'custom-identifier',
        source: LogSources.API,
      }).then((r) => r.unwrap())

      // Wait for the lastResponse promise to resolve
      await lastResponse

      expect(createDocumentLogSpy).toHaveResolvedWith(expect.any(Ok))
      expect(
        createDocumentLogSpy.mock.calls[
          createDocumentLogSpy.mock.calls.length - 1
        ]![0].data.customIdentifier,
      ).toEqual('custom-identifier')
    })

    it('Does not create a document log when rate limited', async () => {
      const streamAIResponseSpy = vi.spyOn(
        await import('../../lib/streamManager/step/streamAIResponse'),
        'streamAIResponse',
      )

      streamAIResponseSpy.mockRejectedValueOnce(
        new ChainError({
          code: RunErrorCodes.RateLimit,
          message: 'Rate limit exceeded',
        }),
      )

      const { context, workspace, document, commit } = await buildData({
        doc1Content: dummyDoc1Content,
      })
      const parameters = { testParam: 'testValue' }
      const { lastResponse } = await runDocumentAtCommit({
        context,
        workspace,
        document,
        commit,
        parameters,
        customIdentifier: 'custom-identifier',
        source: LogSources.API,
      }).then((r) => r.unwrap())

      // Wait for the lastResponse promise to resolve
      await lastResponse

      expect(createDocumentLogSpy).not.toHaveBeenCalled()

      // Restore the spy
      streamAIResponseSpy.mockRestore()
    })
  })

  describe('with userMessage', () => {
    it('injects userMessage into the prompt', async () => {
      const { context, workspace, document, commit } = await buildData({
        doc1Content: dummyDoc1Content,
      })

      const userMessage = 'Please answer in French'
      const { lastResponse } = await runDocumentAtCommit({
        context,
        workspace,
        document,
        commit,
        parameters: {},
        source: LogSources.API,
        userMessage,
      }).then((r) => r.unwrap())

      await lastResponse // wait for the chain to finish

      // Verify that AI was called and the user message was included in one of the calls
      expect(aiSpy).toHaveBeenCalled()

      // Check if any AI call includes the user message
      const calls = aiSpy.mock.calls
      let foundUserMessage = false

      // eslint-disable-next-line @typescript-eslint/no-explicit-any
      calls.forEach((call: any) => {
        const messages = call[0].messages
        // eslint-disable-next-line @typescript-eslint/no-explicit-any
        const userMessages = messages.filter((msg: any) => msg.role === 'user')
        // eslint-disable-next-line @typescript-eslint/no-explicit-any
        userMessages.forEach((msg: any) => {
          // eslint-disable-next-line @typescript-eslint/no-explicit-any
          if (
            msg.content.some(
              (content: any) =>
                content.text &&
                content.text.includes('Please answer in French'),
            )
          ) {
            foundUserMessage = true
          }
        })
      })

      expect(foundUserMessage).toBe(true)
    })

    it('works without userMessage', async () => {
      const { context, workspace, document, commit } = await buildData({
        doc1Content: dummyDoc1Content,
      })

      const { lastResponse } = await runDocumentAtCommit({
        context,
        workspace,
        document,
        commit,
        parameters: {},
        source: LogSources.API,
        // No userMessage provided
      }).then((r) => r.unwrap())

      await lastResponse // wait for the chain to finish

      // Verify that AI was called but no extra user messages were added
      expect(aiSpy).toHaveBeenCalled()

      // Check that no calls include our test user messages
      const calls = aiSpy.mock.calls
      let foundExtraUserMessage = false

      // eslint-disable-next-line @typescript-eslint/no-explicit-any
      calls.forEach((call: any) => {
        const messages = call[0].messages
        // eslint-disable-next-line @typescript-eslint/no-explicit-any
        const userMessages = messages.filter((msg: any) => msg.role === 'user')
        // eslint-disable-next-line @typescript-eslint/no-explicit-any
        userMessages.forEach((msg: any) => {
          // eslint-disable-next-line @typescript-eslint/no-explicit-any
          if (
            msg.content.some(
              (content: any) =>
                content.text &&
                (content.text.includes('Please answer in French') ||
                  content.text.includes('Please respond in Spanish')),
            )
          ) {
            foundExtraUserMessage = true
          }
        })
      })

      expect(foundExtraUserMessage).toBe(false)
    })
  })
})
